Re-enable split-KV for NVFP4 KV cache in prefill path - #1

FA2DetermineCtaTileQ picked a 1x4 warp layout without accounting for the real
shared-memory budget, yielding max_mma_kv=0 at HEAD_DIM_QK=512 under real GQA
with 1-byte (fp8/NVFP4) KV. Make the warp-layout choice smem-aware, and cleanly
reject KV tiles that do not fit real shared memory instead of failing later.

Signed-off-by: Jetha Chan <jethachan@gmail.com>
(cherry picked from commit 23ce85a)

…ackbits

Fixes a device mismatch when building packed custom masks for bidirectional
(multimodal-prefix) attention.

Signed-off-by: Jetha Chan <jethachan@gmail.com>
(cherry picked from commit ee9301f)

Signed-off-by: Jetha Chan <jethachan@gmail.com>
(cherry picked from commit 131bd2a)

…eads

The VO-split rederive left three dispatch gates keyed on head_dim_vo while
FA2DetermineCtaTileQ selects CTA_TILE_Q from the QK dim. For asymmetric
Gemma-4 heads (qk=512, vo=256) VO-split does not engage (NUM_MMA_D_VO==16),
so the full-D path applies and only CTA_TILE_Q=16 fits the 256-register
o_frag wall. Align all four sites:
- utils.cuh FA2DetermineCtaTileQ: VO>=512 -> {16,32}; QK>=512 & VO<512 -> 16
- prefill.cuh KernelTraits::IsInvalid first clause: key on HEAD_DIM_QK
- batch_prefill_{paged,ragged}_kernel_inst.jinja: instantiate {16,32} when head_dim_qk>=512

Validated NVFP4 paged prefill (sm120, RTX): max_abs_err qk256/vo256=0.0048,
qk128=0.0046, qk512/vo256=0.0047 (seq128) / 0.0044 (seq512) / 0.0042
(seq1024) / 0.0166 causal. No NaN; output sized by VO.

Signed-off-by: Jetha Chan <jethachan@gmail.com>

Address review on flashinfer-ai#3684: the paged-decode output allocation doubled
out_head_dim whenever kv_cache_sf was non-None. Gate it on the KV cache
actually being uint8-packed (the NVFP4 layout that stores VO at half
width), so a stray scale-factor tensor on a non-uint8 cache fails the
shape check instead of silently allocating a mis-sized output.

Signed-off-by: Jetha Chan <jethachan@gmail.com>

NVFP4 paged KV is packed uint8 with a per-16-element FP8 block scale. Split-KV
(flash-decoding) chunks the KV range by kv_chunk_size, which is not aligned to the
16-element scale blocks; a chunk boundary landing mid-block -- plus the small
per-split chunk tripping the 1-byte-KV NUM_MMA_KV tile floor -- corrupts the
dequantized reads. It only surfaces when a short query attends a long KV
(qo_len << kv_len), i.e. decode and prefix-cache extend, so dense full-prefill
tests miss it while prefix caching / long-context decode break.

Gate split-KV off when kv_data_type is NVFP4 (uint8 / float4_e2m1fn_x2) in the
paged and ragged prefill plan(). FP8 / 16-bit KV are unaffected. Verified on
Gemma-4 E2B/E4B: radix-on retrieval cliffs at ~600 tokens of reused-prefix
context without the gate, holds to 1448 (== bf16) with it -- on both sm120 (RTX
PRO 6000) and sm121 (GB10 / DGX Spark).

Follow-up: make the FP4 split path scale-block-aware (16-token-aligned chunk
boundaries + NUM_MMA_KV floor) to restore flash-decoding parallelism for
long-context NVFP4 decode.

Signed-off-by: Jetha Chan <jethachan@gmail.com>

Address CodeRabbit review on flashinfer-ai#3684 (5 findings):

prefill.py:
- Gate NVFP4 out_head_dim doubling on `v.dtype == torch.uint8` (the packed
  storage), not just `kv_cache_sf is not None`, at all three prefill sites.
  A stray scale-factor tensor on a non-uint8 cache no longer silently doubles
  the output width. Mirrors the decode-side guard.
- Ragged custom-mask plan(): move mask_indptr to custom_mask.device before
  segment_packbits, mirroring the paged-path fix. mask_indptr inherits
  qo_indptr's device (often CPU) while custom_mask is on GPU.

prefill.cuh:
- init_rope_freq: bound the fill loop on NUM_MMA_D_QK/2, not NUM_MMA_D_VO/2.
  rope_freq is sized [NUM_MMA_D_QK/2][4] and the rotary appliers index it up
  to NUM_MMA_D_QK/2; the asymmetric VO-split dispatch (qk=512, vo=256) made
  NUM_MMA_D_QK != NUM_MMA_D_VO reachable, leaving the upper half uninitialized
  for in-kernel-RoPE callers.
- page_produce_kv_sf deswizzle branch: guard the unconditional smem store on
  flat_byte < SF_TOTAL_BYTES (the rounded-up NUM_SF_ITERS leaves over-range
  lanes). Live cp_async path was already predicated.
- NUM_MMA_KV occupancy budget (all three dispatchers): add the FP4 K/V
  scale-factor smem term, which scales with CTA_TILE_KV. Prevents the budget
  from over-selecting a tile whose SharedStorage then exceeds the smem limit.

Signed-off-by: Jetha Chan <jethachan@gmail.com>

The fwd entry point passes the caller's uninitialized lse (torch.empty)
straight to run_mha_fwd. The kernel writes `out` for every query row but
does not guarantee writing every (batch, head, seq) entry of lse, so lse
can read back whatever the allocator handed out — observed as flaky NaNs
that pass on a clean allocation and fail when a prior test dirtied the
pages. Zero lse before launch, mirroring the existing seq_len==0 branch.

Found smoke-testing flashinfer-ai#3684 on sm120 (RTX PRO 6000): the nvfp4 sm120
accuracy test's lse NaN-check flaked only when run after the paged-prefill
nvfp4 test.

Signed-off-by: Jetha Chan <jethachan@gmail.com>

The old assertion expected a literal "constexpr bool REQUIRE_FP4_KV_CACHE
= true;" that no code path ever emitted; check the #error guard and
static_assert the config template actually renders for FP4 KV.

Signed-off-by: Jetha Chan <jethachan@gmail.com>

BatchDecodeWithPagedKVCacheRun accepts independent K/V stride arrays
since the paged_kv_t independent-stride constructor was introduced, but
the FA2 decode kernel (decode.cuh) still computes a single
protective_get_kv_offset per row -- from the K strides -- and uses it to
address both k_data and v_data, so asymmetric strides would silently
read V at K offsets.

Restore the host-side per-dimension stride-equality check that guarded
this before, with an error message naming the limitation. Restoring the
check (rather than making the decode kernels V-stride-aware) is the
minimal honest fix: no supported path needs asymmetric decode today --
asymmetric NVFP4 (VO-split) decode rides the prefill wrapper, and
symmetric caches have equal K/V strides by construction. If asymmetric
decode is needed later, decode.cuh should grow separate K/V offsets the
way the prefill path already has.

Addresses review feedback on flashinfer-ai#3684 from @qsang-nv.

Signed-off-by: Jetha Chan <jethachan@gmail.com>

…oducer

page_produce_kv_on_the_fly computed both K and V global-memory offsets
via get_paged_kv_offset_for_logical_row, which called
protective_get_kv_offset -- the K strides -- so the shared-KV
(on-the-fly) path addressed V rows with K offsets once K and V may
carry independent strides. Thread produce_v through
get_paged_kv_offset_for_logical_row and select protective_get_v_offset
vs protective_get_k_offset with if constexpr, mirroring what the
prefetched-offset path already does with its separate K/V offset
arrays.

Addresses review feedback on flashinfer-ai#3684 from @qsang-nv and @lesj0610.

Signed-off-by: Jetha Chan <jethachan@gmail.com>

Revert the extension of the b12x auto-preference from SM120-only to all
SM12x, restoring the upstream heuristic: SM121 (GB10) intentionally
keeps cutlass/cudnn first because they are faster there in most cases,
while b12x remains available on SM121 as an explicit backend. If
flipping the SM121 default to b12x is warranted, it should come as a
separately justified PR with benchmarks rather than ride along here.
Also removes the heuristic test that asserted the flipped ordering.

Addresses review feedback on flashinfer-ai#3684 from @qsang-nv.

Signed-off-by: Jetha Chan <jethachan@gmail.com>

Three corrections to the shared-memory feasibility probe in
FA2DetermineCtaTileQ:

* Probe cudaDevAttrMaxSharedMemoryPerBlockOptin instead of the per-SM
  attribute. The kernel dispatcher bounds max_smem_per_threadblock by
  the per-block opt-in limit, so the planner now checks the same limit
  the dispatch-side "even the smallest KV tile exceeds shared memory"
  guard enforces. On parts where the two attributes differ (per-SM
  102400 vs opt-in 101376 on SM86/89/120-class devices), configurations
  in the gap -- e.g. (qk, vo) = (432, 256) at 2-byte KV, 101888 bytes --
  previously passed the probe and then failed dispatch; they now take
  the CTA64 fallback instead.

* Size the KV step by the actual KV dtype width instead of assuming
  2 bytes. sizeof(DTypeKV) is threaded from the batch-prefill plan and
  workspace-size entry points (where DTypeKV is in scope) through
  PrefillPlan / PrefillPlanWorkspaceSize / PrefillPlanImpl /
  PrefillSplitQOKVIndptr as a defaulted trailing parameter, and passed
  directly in single-prefill dispatch; callers that do not know the KV
  dtype keep the previous conservative 2-byte assumption.

* Correct the comment. The old (512, 256) example was unreachable here
  (head_dim_qk >= 512 returns CTA16 before the probe), but the probe is
  not merely forward-looking: plan()/JIT do not validate head dims, so
  within this branch head_dim_qk may be any multiple of 16 below 512
  under pos_encoding_mode NONE, and (qk, vo) = (448, 256) at 2-byte KV
  (104448 bytes) exceeds the 101376-byte opt-in limit of 99KB parts --
  the probe fires today and the CTA64 fallback keeps the configuration
  dispatchable. For the same dims at 1-byte KV the corrected sizing
  flips the selection (59392 bytes fits, so CTA16 replaces the CTA64
  the 2-byte assumption forced): the kv_dtype_bytes change alters tile
  selection for reachable configurations rather than being a no-op on
  current hardware. Both behaviors are pinned by
  test_batch_prefill_paged_cta_tile_q_smem_probe_qk448_vo256.

Addresses review feedback on flashinfer-ai#3684 from @qsang-nv, including the
round-2 correction of this commit's earlier "no reachable configuration
triggers the fallback" claim, which was wrong for the reasons above.

Signed-off-by: Jetha Chan <jethachan@gmail.com>

…ound

The previous rationale claimed a split-KV chunk boundary can land in
the middle of a 16-element scale block. That mechanism is wrong: NVFP4
scale factors group 16 consecutive head-dim elements of a single token,
while split-KV partitions the token axis, so a split boundary never
slices a scale block. Reword the gate docstring, call-site comments and
test docstring to state what is actually known -- corruption was
observed empirically when qo_len << kv_len (decode / prefix-cache
extend) and disappears with split-KV disabled at no measured decode
throughput cost -- and cite the interaction between small per-split KV
chunks and the 1-byte-KV NUM_MMA_KV tile floor as an unconfirmed
hypothesis rather than fact. Also drops project-specific NOTE tags from
nearby comments. Comment-only; no behavior change.

Addresses review feedback on flashinfer-ai#3684 from @qsang-nv.

Signed-off-by: Jetha Chan <jethachan@gmail.com>

run() fills declared-but-unprovided JIT scalars from a fixed mapping
and raised a bare KeyError when a JIT module declares a scalar the
mapping does not know how to derive. Raise a ValueError naming the
scalar and listing the derivable set instead. Also drop the redundant
max(0, ...) clamp on the provided-scalar count:
prepare_jit_additional_args always returns at least one entry per
declared tensor name, so the excess over the tensor-name count cannot
be negative; a comment records that invariant.

Addresses review feedback on flashinfer-ai#3684 from @qsang-nv.

Signed-off-by: Jetha Chan <jethachan@gmail.com>

…ll-VO

For head_dim_qk >= 512 with head_dim_vo < 512 (asymmetric heads),
FA2DetermineCtaTileQ always returns CTA_TILE_Q=16, so the CTA_TILE_Q=32
instantiation in the paged/ragged kernel-instantiation lists can never
be dispatched; drop it and instantiate only CTA16 for that shape class.
The dispatch macro case-32 arm still references the symbol, which stays
unresolved in the module exactly like the pre-existing never-selected
CTA_TILE_Q=32 of symmetric head_dim < 512 modules (verified: shipped
modules carry it as an undefined, lazily-bound symbol). Also updates
the KernelTraits::IsInvalid comment to describe the three-way
CTA_TILE_Q selection.

Addresses review feedback on flashinfer-ai#3684 from @qsang-nv.

Signed-off-by: Jetha Chan <jethachan@gmail.com>

…ual-stride rejection

Completes the support-or-reject-explicitly contract for unequal K/V
strides in tree: every consumer reachable from the updated entry points
must either support independently-strided K/V pools or reject them
loudly.

Support half: asymmetric (head_dim_qk != head_dim_vo) NVFP4 paged
prefill over (512,256) and (256,128) x page_size {1,16} x num_kv_heads
{2,8}, causal. K/V pools and their scale-factor tensors are separately
allocated with genuinely different stride families; bf16 sources are
quantized with the in-tree NVFP4 KV quantization kernel and the FA2
output is checked against a float32 reference attention computed on
nvfp4_kv_dequantize_paged output, so kernel and reference consume the
exact same quantized bytes (dequantization oracle, not a requantized
approximation).

Reject half: the CUDA-core decode entry point
(BatchDecodeWithPagedKVCacheRun) addresses both K and V through a
single set of (K) strides, so its restored ICHECK must fire on K/V
pools whose stride families differ instead of silently misaddressing V.
A positive control with identically padded (equal-stride,
non-contiguous) pools runs and matches the reference, proving the
negative case fails because of the stride inequality and not the padded
allocation.

Addresses review feedback on flashinfer-ai#3684 from @qsang-nv.

Signed-off-by: Jetha Chan <jethachan@gmail.com>

Regression test for the FA2DetermineCtaTileQ shared-memory probe at
head dims that reach it today: plan()/JIT do not validate head dims, so
(qk, vo) = (448, 256) under pos_encoding_mode NONE is accepted, and at
2-byte KV its short-q 1x4-layout cost (104448 bytes) exceeds the
101376-byte per-block opt-in limit of 99KB parts.

The test computes the expected tile from the device's actual opt-in
limit (so the assertion is exact on every architecture), asserts the
planned cta_tile_q via PrefillPlanInfo (the same technique as
test_fp8_prefill.py), and for 2-byte KV runs the kernel against an
exact float32 reference: on 99KB parts this proves the probe fires and
the CTA64 fallback keeps the configuration dispatchable where the
CTA16 dispatch would exceed the per-block limit, and on larger-smem
parts it proves the CTA16 selection runs. For 1-byte KV the assertion
is plan-level: the FA2 1-byte KV producers require head_dim to be a
multiple of 128 elements (the 128-bit-per-lane load loop steps
NUM_MMA_D by 8, and the k128B swizzle needs an 8-aligned upcast
stride), so no currently-runnable 1-byte configuration reaches the
flipped CTA64->CTA16 region -- the pin locks the documented planner
behavior for when one does.

Addresses review feedback on flashinfer-ai#3684 from @qsang-nv.

Signed-off-by: Jetha Chan <jethachan@gmail.com>

…trides

The asymmetric NVFP4 stride test does not execute the produce_v fix in
page_produce_kv_on_the_fly: that producer runs only under
USE_KV_SHARED_SMEM, which excludes FP4 and requires HEAD_DIM_QK ==
HEAD_DIM_VO, so the NVFP4 asymmetric path takes the prefetched
thr_local_kv_offset_{k,v} arrays instead.

Add the configuration that does execute it: 16-bit KV at
(qk, vo) = (512, 512), where USE_KV_SHARED_SMEM holds for both CTA
tiles the planner can pick (static reasoning from prefill.cuh):
CTA_TILE_Q=16 for short q (NUM_WARPS_KV=4; NUM_MMA_D_VO=32 % 4 == 0)
and CTA_TILE_Q=32 for long q (kLargeHeadWarpSplit: NUM_WARPS_KV=2;
32 % 2 == 0), so USE_VO_SPLIT -- and with fp16's equal head dims,
USE_KV_SHARED_SMEM -- is true either way; the qo_len parametrization
covers both tiles and the kv_layout parametrization covers NHD and
HND. K and V pools are views of differently padded parent tensors
(identical logical shapes, unequal stride families, mirroring the
decode negative test's construction), so
get_paged_kv_offset_for_logical_row<produce_v=true> must route V rows
through the V strides: with the fix reverted, the V loads walk K's
stride family and the output diverges from the exact float32
reference, which is how this test was validated to catch the bug it
pins. The configuration is SM80+, so it runs on the standard CI
runners.

Addresses review feedback on flashinfer-ai#3684 from @qsang-nv.

Signed-off-by: Jetha Chan <jethachan@gmail.com>

The (448, 256) CtaTileQ smem-probe test parametrizes over kv_dtype in
{float16, float8_e4m3fn}. On pre-SM100 GPUs the FP8 (1-byte) parametrization
errors before reaching the tile assertion: _fa2_head_dim_nvcc_flags restricts
non-NVFP4 1-byte large-head modules to major versions [10, 11, 12], so the JIT
spec-gen inside plan() raises "No supported CUDA architectures found for major
versions [10, 11, 12]". skip_if_head_dim_unsupported only gates the 16-bit path,
so it misses this.

Add a dtype-aware skip mirroring the module gate, and narrow the docstring
wording "exact on every architecture" -> "exact on every supported
architecture". The fp16 parametrization is unaffected (2-byte fallback, SM80+).

Addresses review feedback on flashinfer-ai#3684 from @qsang-nv.

Signed-off-by: Jetha Chan <jethachan@gmail.com>

gen_customize_batch_prefill_module now requires the scale-factor tensors
(maybe_k_cache_sf / maybe_v_cache_sf) as additional inputs whenever the KV
dtype resolves to NVFP4, raising ValueError otherwise. The host-side
test_customize_batch_prefill_nvfp4_large_head_uses_prefill_flags still called
the generator with empty additional-tensor lists, so it tripped that
ValueError before reaching either flag assertion and failed on every arch
(it never touches the GPU). Pass the two uint8_t SF tensors, mirroring the
production caller, so generation completes and the assertions run:
_fa2_prefill_head_dim_nvcc_flags emits sm_86 (allow_nvfp4_sm8_large_head), and
the plain _fa2_head_dim_nvcc_flags still restricts to [10,11,12] and raises.

Signed-off-by: Jetha Chan <jethachan@gmail.com>